package dslab.util;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import java.io.*;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

/**
 * extension of Reader that supports AES communication and can be initialized in a secure
 * session dependent context midway
 * simply wraps reader if encryption is not yet initialized
 */
public class EncryptedReader extends BufferedReader {
    private final Cipher cipher;
    private boolean initialized = false;

    /**
     * EncryptedReader constructor
     * @param reader InputStreamReader object to wrap
     * @throws NoSuchPaddingException if padding option is invalid
     * @throws NoSuchAlgorithmException if algorithm cannot be used
     */
    public EncryptedReader(InputStreamReader reader) throws NoSuchPaddingException, NoSuchAlgorithmException {
        super(reader);
        this.cipher = Cipher.getInstance("AES/CTR/NoPadding");
    }

    /**
     * reads a line and throws an exception if an error occurred
     * simply wraps `readLine()` of the inherent BufferedReader object if decryption is uninitialized,
     * and attempts decryption if decryption is initialized
     * @return line that was read, without newline characters
     * @throws IOException if a general I/O exception occurred
     * @throws IllegalBlockSizeException if the decryption block size is not allowed
     * @throws BadPaddingException if the decryption padding setting is not allowed
     */
    public String readUnsafeLine() throws IOException, IllegalBlockSizeException, BadPaddingException {
        String string = super.readLine();
        if (initialized)
            string = CryptoUtil.decryptMessage(string, cipher);
        return string;
    }

    /**
     * reads a line in a wrapped context
     * generally exhibits the same semantics as `readLine()` of the inherent BufferedReader object, but uses
     * decryption if AES decryption was initialized
     * @return line that was read, without newline characters
     * @throws IOException if a general I/O exception occurred
     */
    @Override
    public String readLine() throws IOException {
        String string = super.readLine();
        if (initialized) {
            try {
                string = CryptoUtil.decryptMessage(string, cipher);
            } catch (IllegalBlockSizeException | BadPaddingException exception) {
                throw new IOException(exception);
            }
        }
        return string;
    }

    /**
     * initializes this EncryptedReader to switch over to a mode where communication encrypted via AES is used
     * @param key SecretKey to initialize AES communication with
     * @param iv initialization vector to use
     * @throws InvalidAlgorithmParameterException if the encryption algorithm cannot be used
     * @throws InvalidKeyException if the key is not valid
     */
    public void init(SecretKey key, byte[] iv) throws InvalidAlgorithmParameterException, InvalidKeyException {
        cipher.init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(iv));
        initialized = true;
    }
}
